##########################################################################################

library(data.table)
library(optparse)
library(Mfuzz)
library(parallel)
library(dplyr)
library(limma)
library(RColorBrewer)

##########################################################################################

option_list <- list(
    make_option(c("--data_type"), type = "character"),
    make_option(c("--geneset_type"), type = "character"),
    make_option(c("--cluster_file"), type = "character"),
    make_option(c("--clust"), type = "character"),
    make_option(c("--rsem_file"), type = "character"),
    make_option(c("--motif_file"), type = "character"),
    make_option(c("--geneset_file"), type = "character"),
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
    data_type <- "exp"
    geneset_type <- "known_motif"
    clust <- 12
    cluster_file <- "~/20231121_singleMuti/config/cluster_celltype.csv"
    rsem_file <- "~/20231121_singleMuti/results/qc_atac_v3/germ/GeneExpression.All.rds"
    motif_file <- "~/20231121_singleMuti/results/celltype_plot/mfuzz/cor.motif_atac-rna.tsv"
    out_path <- "~/20231121_singleMuti/results/celltype_plot/mfuzz_all_motif"
    geneset_file <- "~/20231121_singleMuti/config/Human_reported_TF2.csv"
}

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

cluster_file <- opt$cluster_file
data_type <- opt$data_type
rsem_file <- opt$rsem_file
motif_file <- opt$motif_file
out_path <- opt$out_path
clust <- as.numeric(opt$clust)
geneset_file <- opt$geneset_file
geneset_type <- opt$geneset_type

dir.create(out_path , recursive = T)

##########################################################################################

dat_tpm <- readRDS(rsem_file)
dat_motif <- data.frame(fread(motif_file , header = T))
dat_cluster <- data.frame(fread(cluster_file , header = F))

set.seed(1234)

##########################################################################################
## 细胞顺序
col_order <- c("SSC","Differenting&Differented SPG",
    "Leptotene","Zygotene","Patchytene","Diplotene",
    "Early stage of spermatids","Round&ElongateS.tids","Sperm"
    )

## 对应clsuter和细胞类型
dat_cluster$V1 <- gsub( "cluster" , "" , dat_cluster$V1 )
dat_cluster$V2 <- gsub( " " , "\n" , dat_cluster$V2 )
dat_cluster$V2 <- gsub( "&" , "\n" , dat_cluster$V2 )

col_order <- gsub( " " , "\n" , col_order )
col_order <- gsub( "&" , "\n" , col_order )

##########################################################################################
## 提取感兴趣的基因集合
if(geneset_type == "all_motif"){
    dat_geneset <- data.frame(fread(geneset_file , header = T))
    tf_name <- dat_geneset$MotifMatrix_matchName
}else if(geneset_type == "known_motif"){
    dat_geneset <- data.frame(fread(geneset_file , header = F))
    tf_name <- dat_geneset$V1
    tf_name <- tf_name[tf_name %in% rownames(dat_tpm)]
}

colnames(dat_tpm) <- gsub( "#" , "_" , colnames(dat_tpm) )
dat_tpm <- dat_tpm[tf_name,]

##########################################################################################
## 画图
## 标记每个cluster对应多少基因
mfuzz.plot_revise <- function (eset, cl, mfrow = c(1, 1), colo, min.mem = 0, time.labels, 
    new.window = TRUE){

    clusterindex <- cl[[3]]
    memship <- cl[[4]]
    memship[memship < min.mem] <- -1
    colorindex <- integer(dim(exprs(eset))[[1]])

    if (missing(colo)) {
        colo <- c("#FF8F00", "#FFA700", "#FFBF00", "#FFD700", 
            "#FFEF00", "#F7FF00", "#DFFF00", "#C7FF00", "#AFFF00", 
            "#97FF00", "#80FF00", "#68FF00", "#50FF00", "#38FF00", 
            "#20FF00", "#08FF00", "#00FF10", "#00FF28", "#00FF40", 
            "#00FF58", "#00FF70", "#00FF87", "#00FF9F", "#00FFB7", 
            "#00FFCF", "#00FFE7", "#00FFFF", "#00E7FF", "#00CFFF", 
            "#00B7FF", "#009FFF", "#0087FF", "#0070FF", "#0058FF", 
            "#0040FF", "#0028FF", "#0010FF", "#0800FF", "#2000FF", 
            "#3800FF", "#5000FF", "#6800FF", "#8000FF", "#9700FF", 
            "#AF00FF", "#C700FF", "#DF00FF", "#F700FF", "#FF00EF", 
            "#FF00D7", "#FF00BF", "#FF00A7", "#FF008F", "#FF0078", 
            "#FF0060", "#FF0048", "#FF0030", "#FF0018")
    }
    colorseq <- seq(0, 1, length = length(colo))
    for (j in 1:max(clusterindex)) {
        tmp <- exprs(eset)[clusterindex == j, , drop = FALSE]
        tmpmem <- memship[clusterindex == j, j]
        if (((j - 1)%%(mfrow[1] * mfrow[2])) == 0) {
            if (new.window) 
                X11()
            par(mfrow = mfrow)
            if (sum(clusterindex == j) == 0) {
                ymin <- -1
                ymax <- +1
            }
            else {
                ymin <- min(tmp)
                ymax <- max(tmp)
            }
            plot.default(x = NA, xlim = c(1, dim(exprs(eset))[[2]]), 
                ylim = c(ymin, ymax), xlab = "", ylab = "Expression changes", 
                main = paste("Cluster", j , "(" , nrow(tmp) , ")"), axes = FALSE)
            if (missing(time.labels)) {
                axis(1, 1:dim(exprs(eset))[[2]], c(1:dim(exprs(eset))[[2]]))
                axis(2)
            }
            else {
                axis(1, 1:dim(exprs(eset))[[2]], time.labels)
                axis(2)
            }
        }
        else {
            if (sum(clusterindex == j) == 0) {
                ymin <- -1
                ymax <- +1
            }
            else {
                ymin <- min(tmp)
                ymax <- max(tmp)
            }
            plot.default(x = NA, xlim = c(1, dim(exprs(eset))[[2]]), 
                ylim = c(ymin, ymax), xlab = "", ylab = "Expression changes", 
                main = paste("Cluster", j , "(" , nrow(tmp) , ")"), axes = FALSE)
            if (missing(time.labels)) {
                axis(1, 1:dim(exprs(eset))[[2]], c(1:dim(exprs(eset))[[2]]))
                axis(2)
            }
            else {
                axis(1, 1:dim(exprs(eset))[[2]], time.labels)
                axis(2)
            }
        }
        if (!(sum(clusterindex == j) == 0)) {
            for (jj in 1:(length(colorseq) - 1)) {
                tmpcol <- (tmpmem >= colorseq[jj] & tmpmem <= 
                  colorseq[jj + 1])
                if (sum(tmpcol) > 0) {
                  tmpind <- which(tmpcol)
                  for (k in 1:length(tmpind)) {
                    lines(tmp[tmpind[k], ], col = colo[jj])
                  }
                }
            }
        }
    }
}

##########################################################################################
## 按时间表达聚类
#clust <- 12

DEGs_exp <- dat_tpm
time <- data.frame( 
    sample = colnames(DEGs_exp) ,
    class = sapply( strsplit(colnames(DEGs_exp) , "_" ) , "[" , 3 ) 
)

time <- merge( time , dat_cluster , by.x = "class" , by.y = "V1" )
time <- time[,c(2,3)]
colnames(time) <- c( "sample" , "class" )

tmp <- data.frame(colnames(DEGs_exp),t(DEGs_exp))
temp <- data.frame(time[match(tmp[,1],time[,1]),],tmp)
DEGs_exp_averp <- t(limma::avereps(temp[,-c(1:3)],ID=temp[,2]))
DEGs_exp_averp <- DEGs_exp_averp[,col_order]

## 1. 预处理：去除表达量太低或者在不同时间点间变化太小的基因等步骤
# Mfuzz聚类时要求是一个ExpressionSet类型的对象，所以需要先用表达量构建这样一个对象。
eset <- new("ExpressionSet",exprs = DEGs_exp_averp)

# 根据标准差去除样本间差异太小的基因
eset <- filter.std(eset,min.std=0)

## 2. 标准化：聚类时需要用一个数值来表征不同基因间的距离，Mfuzz中采用的是欧式距离，
# 由于普通欧式距离的定义没有考虑不同维度间量纲的不同，所以需要先进行标准化
eset <- standardise(eset)

## 3. 聚类：Mfuzz中的聚类算法需要提供两个参数，第一个参数为希望最终得到的聚类的个数，这个参数由我们直接指定；
# 第二个参数称之为fuzzifier值，用小写字母m表示，可以通过函数评估一个最佳取值
#clust <- 9 # 聚类个数，文章中用的6个聚类，我们也用6个
m <- mestimate(eset) #  评估出最佳的m值
cl <- mfuzz(eset, c = clust, m = m) # 聚类

out_name <- paste0(out_path , "/mfuzz_plot." , data_type , "." , clust , ".pdf")
color.2 <- colorRampPalette(rev(c("#ff0000", "Yellow", "OliveDrab1")))(1000)
pdf(out_name , width = 15)
mfuzz.plot_revise(eset,cl,mfrow=c(1,1),new.window= FALSE,time.labels=colnames(DEGs_exp_averp),colo = color.2)
dev.off()

# 在cl这个对象中就保存了聚类的完整结果，对于这个对象的常见操作如下
# cl$size # 查看每个cluster中的基因个数
# [1] 2269 1982 2289 1653 1393 1729

# cl$cluster[cl$cluster == 1] # 提取某个cluster下的基因
# cl$membership # 查看基因和cluster之间的membership

## acore函数帮助我们计算出来了每个趋势分组里面的每个基因的重要性
acore <- acore(eset , cl , min.acore=0)
acore_list <- do.call(rbind, lapply(seq_along(acore), function(i){ data.frame(CLUSTER=i, acore[[i]])}))
## 标记是否为TF
acore_list$TF <- ifelse(acore_list$NAME %in% tf_name , TRUE , FALSE)
DEGs_exp_averp <- data.frame(DEGs_exp_averp)
colnames(DEGs_exp_averp) <- gsub( "\n" , " " ,  colnames(DEGs_exp_averp)  )
DEGs_exp_averp$NAME <- rownames(DEGs_exp_averp)

acore_list <- merge( acore_list , DEGs_exp_averp , by = "NAME" )
acore_list <- merge( acore_list , dat_motif , by.x = "NAME" , by.y = "MotifMatrix_matchName" )

out_name <- paste0(out_path , "/mfuzz_plot." , data_type , "." , clust , ".tsv")
write.table( acore_list , out_name , row.names = F , quote = F , sep = "\t" )
